# -*- coding: utf-8 -*-
# @Time    : 2023/5/12 9:13 下午
# @Author  : Wu WanJie

import os
import pandas as pd
import warnings
import seaborn as sns
import matplotlib.pyplot as plt

warnings.filterwarnings("ignore")


def load_csv_dataset(csv_path):
    df = pd.read_csv(csv_path)
    text_lst = []
    labels_lst = []

    for row in df.itertuples():
        text_lst.append(getattr(row, "text"))
        tmp_lst = []
        for _label in getattr(row, "label").split(";"):
            if "-" in _label:
                key = _label.split("-")[0]
            elif "（" in _label:
                key = _label.split("（")[0]
            else:
                key = _label

            tmp_lst.append(key)
        labels_lst.append(tmp_lst)
    return {
        "text_lst": text_lst,
        "labels_lst": labels_lst
    }


def get_new_data_frame(data, new_path):
    labels = set()
    labels_lst = data["labels_lst"]
    [[labels.add(j) for j in lst] for lst in labels_lst]
    new_tas = {label: list() for label in labels}
    for tmp_labels in labels_lst:
        for label in labels:
            if label in tmp_labels:
                new_tas[label].append(1)
            else:
                new_tas[label].append(0)

    new_tas["text"] = data["text_lst"]
    df = pd.DataFrame(new_tas)
    df.to_csv(new_path, index=False)


def convert_new_csv_format(data_dir):
    train_data = load_csv_dataset(os.path.join(data_dir, "train_data.csv"))
    get_new_data_frame(train_data, os.path.join(data_dir, "train_data.csv"))

    eval_data = load_csv_dataset(os.path.join(data_dir, "eval_data.csv"))
    get_new_data_frame(eval_data, os.path.join(data_dir, "eval_data.csv"))


def label_distribute(new_train_df):
    total_columns = new_train_df.columns.tolist()
    total_columns.remove("text")
    sns.set(
        style='whitegrid',  # 设置背景风格为白底加网格线，使得图表更加清晰
        palette='muted',  # 设置颜色调色板为淡色调，以减少图表的突兀度
        font_scale=0.6  # 设置字体大小的缩放比例，以便更好地适应图表的大小
    )
    plt.rcParams["font.family"] = ["SimHei"]
    plt.rcParams["axes.unicode_minus"] = False

    new_train_df[total_columns].sum().sort_values().plot(kind="barh")
    plt.show()


def text_distribute(new_train_df_path):
    new_train_df = os.path.join(new_train_df_path, "train_data.csv")
    text_lst = list()
    text_len = list()
    df = pd.read_csv(new_train_df)
    for row in df.itertuples():
        text_lst.append(getattr(row, "text"))
        text_len.append(len(getattr(row, "text")))
    result = pd.DataFrame({"text_lst": text_lst, "text_len": text_len})
    print(result["text_len"].describe())
    print(result["text_lst"].describe())


if __name__ == "__main__":
    data_dir_out = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "data/competition")
    convert_new_csv_format(data_dir_out)

    # ===============统计标签分布情况=====================
    new_train_df_out = pd.read_csv(os.path.join(data_dir_out, "train_data.csv"))
    label_distribute(new_train_df_out)

    # ===============统计文本长度分布情况=====================
    text_distribute(data_dir_out)
